#!/usr/bin/env python

# --------
# Imports:
# --------
# Project configuration
from config import *

# Ros imports:
import roslib; roslib.load_manifest(NODE_NAME)
import rospy
from tf import TransformListener, Exception as TFException
from geometry_msgs.msg import Twist, PointStamped
from sensor_msgs.msg import LaserScan
from nav_msgs.msg import Odometry

# Other imports:
import os, threading, time
from Queue import Queue
from copy import deepcopy
from numpy.random import normal

# Our project
from particle import Particle
from slam_map import SlamMap
from ancestry_tree import AncestryTree
from misc import *
from wu_dpslam.srv import *
from localization_thread import *
from map_update_thread import *

# ---------
# Constants:
# ---------
UNIFORM_PARTICLE_PROBABILITY = 1.0 / NUM_PARTICLES

# ----------------------------
class DPSlam(threading.Thread):
# ----------------------------
	# -----------------
	def __init__(self):
	# -----------------
		# Initialize ourselves as a thread
		threading.Thread.__init__(self)

		# Core datastructures
		# -------------------
		self.ancestryTree = None						# Essentially stores the hierarchy of particles in consideration
		self.slamMap = None								# Stores the map data structure
		self.localizationJobQueue = Queue()				# A thread safe container to store localization "jobs" (particle, laser)
		self.localizationThreads = []					# A list of localization worker threads
		self.mapUpdateJobQueue = Queue()				# A thread safe container to store map update "jobs" (particle, laser)
		self.mapUpdateThreads = []						# A list of map update worker threads
	
		# Stored messages & relevant info
		# -------------------------------
		self.running = False 							# True if the DPSLAM thread is running
		self.laserMsg = None							# The most recent laser scan message we received
		self.odomMsg = None								# The most recent odometry message we received
		self.transformListener = None					# Used to access tf frames

		# Service handlers
		# ----------------
		self.occupancyGridService = rospy.Service(OCCUPANCY_GRID_SERVICE_NAME, OccupancyGridService, self.handleGetOccupancyGrid)
		self.ancestryTreeVisualizationService= rospy.Service(ANCESTRY_TREE_VISUALIZATION_SERVICE_NAME, AncestryTreeVisualizationService, self.handleGetAncestryTreeVisualization)
	
		# Subscriptions
		rospy.Subscriber(LASER_SCAN_TOPIC, LaserScan, self.laserCallback)
		rospy.Subscriber(ODOMETRY_TOPIC, Odometry, self.odometryCallback)
		self.transformListener = TransformListener()
	
	# ----
	# Core|
	# -----------
	def run(self):
	# -----------
		# Print out some status information
		rospy.loginfo(' -------------')
		rospy.loginfo(' D P - S L A M:')
		rospy.loginfo(' -------------')

		# We're running now
		self.running = True

		# Block until we recieve an odometry message and a laser scan message
		rospy.loginfo('+ Blocking for an initial odometry message and laser scan message')
		laser, position, heading = self.blockForMessages()
		prevPosition = position
		prevHeading = heading

		if laser is None or position is None: self.stop()
		if not self.running: return
	
		# Create the map data structure
		self.slamMap = SlamMap(MAP_WIDTH_PX, MAP_HEIGHT_PX, MAP_METERS_PER_PIXEL, MAP_CENTER_PX, position, heading)
	
		# Create the root particle (no parent, centered at the robot's starting position)
		rospy.loginfo('+ Initialize the root node')
		root = Particle(None, self.slamMap)
		root.position = (0,0)
		root.heading = 0
		rospy.loginfo('  - Initializing root node position: (%f, %f)', position[0], position[1])
		rospy.loginfo('  - Initializing root node heading: %f radians',  (heading))

		# Initialize the root node's map (initialize the sensor model)
		rospy.loginfo('  - Initializing root node map')
		root.updateMap(laser)

		# Create the ancestry tree data structure (giving it the root particle to start)
		rospy.loginfo('+ Initializing the ancestry tree')
		self.ancestryTree = AncestryTree(NUM_PARTICLES, root)

		# Spawn localization worker threads
		rospy.loginfo('+ Spawning %i localization threads', NUM_LOC_THREADS)
		self.localizationThreads = [LocalizationThread(i, self.localizationJobQueue, self) for i in range(NUM_LOC_THREADS)]
		for locThread in self.localizationThreads: 
			rospy.loginfo('  - Started thread #%i', locThread.idNum)
			locThread.start()

		# Spawn map update worker threads
		rospy.loginfo('+ Spawning %i map update threads', NUM_MAP_UPDATE_THREADS)
		self.mapUpdateThreads = [MapUpdateThread(i, self.mapUpdateJobQueue, self) for i in range(NUM_MAP_UPDATE_THREADS)]
		for mapUpdateThread in self.mapUpdateThreads:
			rospy.loginfo('  - Started thread #%i', mapUpdateThread.idNum)
			mapUpdateThread.start()


		# Main Control Loop:
		# -----------------
		# Log status info
		rospy.loginfo('')
		rospy.loginfo('  -------------------------- ')
		rospy.loginfo('+ Starting main control loop:')
		rospy.loginfo('  -------------------------- ')

		# Start the control loop
		iteration = 0
		while self.running:

			# Status info
			rospy.loginfo('')
			rospy.loginfo('+ Iteration #%i:' % iteration)
			iteration += 1

			# ----------------
			# Debugging Output:
			# ----------------
			if RECORD_IMAGES and iteration % ITERATIONS_PER_IMAGE == 0:
				rospy.loginfo('  - Saving image')

				# Find the highest probability leaf
				highestProbLeaf = self.ancestryTree.findHighestProbLeaf()
				
				# DEBUGGING: Save the map of the highest probability leaf
				img = highestProbLeaf.generateVisualization()
				img.save('vis/map%i.png' % (iteration - 1), 'PNG')


			# Block for odometry and a laser information
			rospy.loginfo('  - Blocking for a laser reading and an odometry message')	
			laser, position, heading = self.blockForMessages()

			if laser is None or position is None or heading is None: self.stop()
			if not self.running: break
			
			# ------------
			# Motion Model:
			# ------------
			# Calculate the change in position
			deltaX = position[0] - prevPosition[0]
			deltaY = position[1] - prevPosition[1]
			deltaH = clipAngle(heading - prevHeading)
			deltaOdom = (deltaX, deltaY, deltaH)

			# Update the previous odometry message
			prevPosition = position
			prevHeading = heading

			# For each particle, run the motion model
			rospy.loginfo('  - Running the motion model for each particle leaf')
			for particle in self.ancestryTree:
				particle.motionModel(deltaOdom)

			# Make sure we're still running
			if not self.running: break

			# ------------
			# Localization:
			# ------------
			# Build a (synchronized) queue of localization 'jobs' for the worker threads to execute
			localizationStart = time.time()
			rospy.loginfo('  - Running Localization')
			rospy.loginfo('    * Adding localization jobs to queue')

			for particle in self.ancestryTree:
				self.localizationJobQueue.put((particle, laser))

			# Block until all the of the localization worker threads are done (this is admittedly a little ugly)
			rospy.loginfo('    * Waiting for localization threads to finish')
			while self.running:	
				breakLoop = True
				for locThread in self.localizationThreads:
					if locThread.running: continue
				if self.localizationJobQueue.empty():
					break

			rospy.loginfo('     * Total localization duration: %fs', time.time() - localizationStart) 

			# Normalize the particles' probabilities
			self.ancestryTree.normalizeProbabilities(self.ancestryTree.particleLeafs)
			self.ancestryTree.sortParticles(self.ancestryTree.particleLeafs)

			# -------
			# Mapping:
			# -------
			# Update the maps of the particles likely to be resampled
			mappingStart = time.time()
			rospy.loginfo('  - Updating the map for high probability particle leafs')
			rospy.loginfo('    * Adding map update jobs to the queue')
			for i, particle in enumerate(self.ancestryTree):
				if i > MAX_NUM_MAP_UPDATES or not self.running: break

				# Only update the map if the particle is likely to be resampled
				if particle.probability >= UNIFORM_PARTICLE_PROBABILITY:
					#particle.updateMap(laser)
					self.mapUpdateJobQueue.put((particle, laser))
				
				# If we're not updating the map, ensure that the particle won't be sampled again
				else:
					particle.probability = 0
		
			# Block until all the of the map update worker threads are done (this is again admittedly a little ugly)
			rospy.loginfo('    * Waiting for map update threads to finish')
			while self.running:	
				breakLoop = True
				for mapUpdateThread in self.mapUpdateThreads:
					if mapUpdateThread.running: continue
				if self.mapUpdateJobQueue.empty():
					break

			rospy.loginfo('    * Total duration: %fs', time.time() - mappingStart)

			# Make sure we're still running
			if not self.running: break

			# ----------------------------------
			# Importance Sampling & Housekeeping: 
			# ----------------------------------
			# Run importance sampling
			rospy.loginfo('  - Running importance sampling')
			self.ancestryTree.importanceSampling()
		
		# Final status info message	
		rospy.loginfo('+ DP-SLAM Stopped.')

	# ------------
	def stop(self):
	# ------------
		rospy.loginfo('')
		rospy.loginfo('Stopping DP-SLAM...')
		rospy.loginfo('')
		self.running = False

	# ------------------------
	def blockForMessages(self):
	# ------------------------
		# Initializations
		laserScan = None
		position = None
		heading = None

		blockStart = time.time()
		while laserScan is None or position is None or heading is None:
			if time.time() - blockStart < MSG_BLOCK_DURATION_S:
				
				# Make sure we're still running
				if not self.running: return (None, None, None)

				try:
					# Laser
					laserScan = self.laserMsg

					# Position/heading via odom message
					if self.odomMsg is not None:
						odomMsg = self.odomMsg
						position, heading = ((odomMsg.pose.pose.position.x, odomMsg.pose.pose.position.y), quat2heading(odomMsg.pose.pose.orientation))

					# Position/heading via tf
					elif self.transformListener.frameExists(ODOM_FRAME) and self.transformListener.frameExists(BASE_FRAME):
						latestCommonTime = self.transformListener.getLatestCommonTime(ODOM_FRAME, BASE_FRAME)
						position, orientation = self.transformListener.lookupTransform(ODOM_FRAME, BASE_FRAME, latestCommonTime)
						heading = quat2heading(orientation)
				except TFException as e:
					continue
			else:
				rospy.loginfo('Laser: %s, Position: %s, Heading: %s', str(laserScan is not None), str(position is not None), str(heading is not None))
				return (None, None, None)

		# Null out the stured messages	
		self.laserScan = None
		self.odomMsg = None			

		return (laserScan, deepcopy(position), clipAngle(heading))
	
	# ---------
	# Callbacks|
	# --------------------------
	def laserCallback(self, msg):
	# --------------------------
		# Store the laser message
		self.laserMsg = msg

	# -----------------------------
	def odometryCallback(self, msg):
	# -----------------------------
		# Store the odometry message
		self.odomMsg = msg
	
	# --------
	# Services|	
	# -----------------------------------
	def handleGetOccupancyGrid(self, req):
	# -----------------------------------
		rospy.loginfo('*SERVICE* ---> Recieved request for occupancy grid')
		if self.ancestryTree.highestProbLeaf:
			occupancyGrid = self.ancestryTree.highestProbLeaf.generateOccupancyGrid()	
			rospy.loginfo('*SERVICE* ---> Occupancy grid generated')
			return OccupancyGridServiceResponse(occupancyGrid)
	
	# -----------------------------------------------
	def handleGetAncestryTreeVisualization(self, req):
	# -----------------------------------------------
		rospy.loginfo('*SERVICE* ---> Recieved request for ancestry tree visualization')
		return AncestryTreeVisualizationServiceResponse()

# ------------------------
if __name__ == '__main__':
# ------------------------
	# Initialize/register the wu_dpslam node
	rospy.init_node(NODE_NAME)

	# Create our controller
	controller = DPSlam()

	# Run the controller (in its own thread)
	controller.start()

	# Keep the main thread running as long as the controller is running
	rospy.on_shutdown(controller.stop)
	rospy.spin()
